import tensorflow as tf
import numpy as np
import pickle

with open("data_set.pickle", 'rb') as f:
    data_set = pickle.load(f)
x_train = data_set["x_train"]
y_train = data_set["y_train"]
x_test = data_set["x_test"]
y_test = data_set["y_test"]
print(x_train.shape)
print(y_train.shape)
print(x_test.shape)
print(y_test.shape)
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=[56, 56]),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])
# 编译
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
# 训练模型
history = model.fit(x_train,
                    y_train,
                    epochs=5,
                    batch_size=64,
                    validation_data=(x_test, y_test))
np.save("normal.npy", history.history)
